package com.jd.jobkeeper.client;

import com.jd.jobkeeper.protocals.ErrorCode;
import com.jd.jobkeeper.protocals.Response;
import com.jd.jobkeeper.protocals.ResponseStatus;
import com.jd.jobkeeper.protocals.Task;
import com.jd.jobkeeper.transport.Transport;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Scheduler;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

public class JobClientProtocolImpl implements JobClientProtocol {

    private final Logger log = LoggerFactory.getLogger(JobClientProtocol.class);
    private final Map<String, Handler> handlerRegistry = new ConcurrentHashMap<>();

    final Scheduler scheduler;

    final Sinks.Many<Task> finishedMany;
    final ConcurrentMap<Long, TaskContext> taskContexts = new ConcurrentHashMap<>();
    final AtomicInteger counter = new AtomicInteger();
    private Transport transport;

    public JobClientProtocolImpl(Scheduler scheduler,
                                 Sinks.Many<Task> finishedMany) {
        this.scheduler = scheduler;
        this.finishedMany = finishedMany;
    }

    @Override
    public Mono<Void> run(Task task, Transport transport) {

        this.transport.send("testAndSet", task, Response.class)
                .publishOn(scheduler)
                .subscribe(r -> {

                            TaskContext context = new TaskContext(task);
                            Handler handler = handlerRegistry.get(task.getProcessorInfo());
                            if (handler == null) {
                                /*return Mono.just(buildResponse(ResponseStatus.ERROR, ErrorCode.SERVER_UNAVAILABLE));*/
                            }

                            try {


                                TaskContext prev = taskContexts.putIfAbsent(task.getId(), context);
                                if (prev != null) {
                                    counter.decrementAndGet();
                                    /* return Mono.just(buildResponse(ResponseStatus.OK, ErrorCode.TASK_RUNNING));*/
                                }
                                //Test and set task
                                handler.handler(context);
                            } catch (Throwable e) {
                                log.error("Handler task exception : " + task, e);
                            } finally {
                                this.reportTask(context);
                                counter.decrementAndGet();
                                taskContexts.remove(task.getId());
                            }
                        },
                        e -> {
                        });


        return Mono.empty();
    }

    @Override
    public Mono<Response> kill(Task task) {
        TaskContext context = taskContexts.get(task.getId());
        if (context != null) {
            context.interrupt();
        }
        return Mono.just(buildResponse(ResponseStatus.OK, ErrorCode.TASK_RUNNING));
    }


    @Override
    public void interrupt() {
        taskContexts.forEach((k, ctx) -> ctx.interrupt());
    }

    @Override
    public Mono<Response> isRunning(Task task) {
        boolean isRunning = taskContexts.containsKey(task.getId());
        ErrorCode errorCode = isRunning ? ErrorCode.TASK_RUNNING : ErrorCode.SERVER_UNAVAILABLE;
        return Mono.just(buildResponse(ResponseStatus.OK, errorCode));
    }

    public void stop() {
        this.interrupt();
        finishedMany.tryEmitComplete();
        this.scheduler.dispose();
        taskContexts.clear();
        counter.set(0);
        this.transport.dispose();
    }

    private void reportTask(TaskContext context) {
        if (this.transport == null) {
            log.warn("The transport is closed , current task is stop!");
            return;
        }

        if (context.isInterrupted()) {
            log.warn("The task is interrupted!");
            //do not report
            return;
        }

        Task t = context.task;

        this.transport.send("reportTask", t, Response.class)
                .subscribe(response -> {
                        },
                        e -> {
                            log.error("Error report task task.id:" + t.getId(), e);
                        });
    }


    @Override
    public void onConnected(Transport t) {
        this.transport = t;
        this.transport.onClose().doFinally(s -> {
            log.info("The transport is close ....");
            this.transport = null;
            interrupt();
        }).subscribe();
    }

    @Override
    public void register(Handler handler) {
        handlerRegistry.put(handler.name(), handler);
    }

    private Response okResponse() {
        return Response.newBuilder()
                .setStatus(ResponseStatus.OK)
                .build();
    }

    private Response buildResponse(ResponseStatus status, ErrorCode errorCode) {
        return Response.newBuilder()
                .setStatus(status)
                .setError(errorCode)
                .build();
    }
}
